#include <iostream>
#include <fstream>
#include <chrono>
#include <immintrin.h>  // AVX指令集

// 内存对齐分配器
template<typename T>
T* aligned_alloc(size_t size, size_t alignment = 32) {
    void* ptr = nullptr;
    posix_memalign(&ptr, alignment, size * sizeof(T));
    return reinterpret_cast<T*>(ptr);
}

// Sobel算子内核
const int8_t GX[3][3] = {{-1, 0, 1}, {-2, 0, 2}, {-1, 0, 1}};
const int8_t GY[3][3] = {{-1, -2, -1}, {0, 0, 0}, {1, 2, 1}};

// 基础版本：朴素实现
void sobel_baseline(const uint8_t* src, float* dst, int w, int h) {
    for (int y = 1; y < h-1; ++y) {
        for (int x = 1; x < w-1; ++x) {
            int sum_x = 0, sum_y = 0;
            // 卷积计算
            for (int ky = -1; ky <= 1; ++ky) {
                for (int kx = -1; kx <= 1; ++kx) {
                    int offset = (y+ky)*w + (x+kx);
                    sum_x += src[offset] * GX[ky+1][kx+1];
                    sum_y += src[offset] * GY[ky+1][kx+1];
                }
            }
            dst[y*w + x] = sqrtf(sum_x*sum_x + sum_y*sum_y);
        }
    }
}

// SIMD优化版本（处理8像素/周期）
void sobel_simd(const uint8_t* src, float* dst, int w, int h) {
    // 加载水平核系数
    __m256i gx0 = _mm256_setr_epi8(-1,0,1, -2,0,2, -1,0,1, 
                                  -1,0,1, -2,0,2, -1,0,1);
    // 类似加载垂直核...
    
    for (int y = 1; y < h-1; ++y) {
        const uint8_t* row0 = &src[(y-1)*w];
        const uint8_t* row1 = &src[y*w];
        const uint8_t* row2 = &src[(y+1)*w];
        
        for (int x = 1; x < w-1; x += 8) {
            // 加载三行数据
            __m128i r0 = _mm_loadu_si128((__m128i*)(row0 + x - 1));
            __m128i r1 = _mm_loadu_si128((__m128i*)(row1 + x - 1));
            __m128i r2 = _mm_loadu_si128((__m128i*)(row2 + x - 1));
            
            // 转换为16位整型
            __m256i v0 = _mm256_cvtepu8_epi16(r0);
            __m256i v1 = _mm256_cvtepu8_epi16(r1);
            __m256i v2 = _mm256_cvtepu8_epi16(r2);
            
            // 水平梯度计算
            __m256i gx = _mm256_maddubs_epi16(v0, gx0);
            gx = _mm256_add_epi16(gx, _mm256_maddubs_epi16(v1, gx1));
            gx = _mm256_add_epi16(gx, _mm256_maddubs_epi16(v2, gx2));
            
            // 垂直梯度计算（类似逻辑）
            // ...
            
            // 计算幅值并存储
            __m256 f_gx = _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_extracti128_si256(gx, 0)));
            __m256 f_gy = _mm256_cvtepi32_ps(_mm256_cvtepi16_epi32(_mm256_extracti128_si256(gy, 0)));
            __m256 mag = _mm256_sqrt_ps(_mm256_add_ps(_mm256_mul_ps(f_gx, f_gx), 
                                       _mm256_mul_ps(f_gy, f_gy)));
            _mm256_storeu_ps(&dst[y*w + x], mag);
        }
    }
}

int main() {
    const int W = 512, H = 512;
    
    // 加载输入图像
    uint8_t* input = aligned_alloc<uint8_t>(W*H);
    std::ifstream fin("input.bin", std::ios::binary);
    fin.read(reinterpret_cast<char*>(input), W*H);
    
    // 初始化输出
    float* output = aligned_alloc<float>(W*H);
    memset(output, 0, W*H*sizeof(float));
    
    // 性能测试
    auto start = std::chrono::high_resolution_clock::now();
    sobel_baseline(input, output, W, H);
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "Baseline time: " 
              << std::chrono::duration<double>(end - start).count() << "s\n";
    
    // 保存结果
    std::ofstream fout("output.bin", std::ios::binary);
    fout.write(reinterpret_cast<char*>(output), W*H*sizeof(float));
    
    free(input);
    free(output);
    return 0;
}